import numpy as np
import cv2,os,sys,time
import cv2.aruco as aruco

def make_folder(name):
  try:
    os.system("mkdir "+name)
  except:
    pass


def getBottomLeftWhite(im):
  (width,height) = im.shape
  for y in range(height-1,-1,-1):
    for x in range(width):
      if(im[y][x] == 255):
        return (x,y)
  return None

def getTag(num,aruco_dict=aruco.getPredefinedDictionary(aruco.DICT_4X4_100),size=500):
  return aruco.generateImageMarker(aruco_dict, num, size)

def save_tags(aruco_dict,name,num,size=500,flip=False,label=False):
  for i in range(num):
    if i%20==0: print("tag %d generated." % (i))
    im = getTag(i,aruco_dict,size)
    if flip:
      im = cv2.flip(im,1)
    pos = getBottomLeftWhite(im)
    pos = (pos[0]+5,pos[1]-5) #shift up a little
    final=cv2.putText(im,str(i),pos,cv2.FONT_HERSHEY_DUPLEX,2,128, 2) #write num in gray

    if label:
      #add label
      final = np.concatenate((final,255*np.ones((int(size/10),size))))
      msg = "Tag %d" % (i)
      pos = (int(size/2)-20*int(len(msg)/2),size+int(size/20))
      final = cv2.putText(final,"Tag %d" % (i),pos,cv2.FONT_HERSHEY_COMPLEX_SMALL,1,0)

    cv2.imwrite(name+str(i)+".jpg",final)


def generate_tags(dict_name,outfilename,quantity=100,flip=False,label=False):
  aruco_dict = aruco.getPredefinedDictionary(dict_name)
  save_tags(aruco_dict,outfilename,quantity,flip=flip,label=label)


if(__name__ == "__main__"):
  print("you are running this file as a standalone program.")
  label = len(sys.argv)>1
  if(label):
    print("You have chosen to label images of all tags.")
  print("tags being outputed will be saved to autogenerated folders in your current directory. Press enter to continue?")
  input() #wait for user to press enter
  make_folder("aruco_4x4_100")
  #make_folder("aruco_4x4_1000")
  #make_folder("aruco_5x5_100")
  #make_folder("aruco_5x5_1000"),
  generate_tags(aruco.DICT_4X4_100,"aruco_4x4_100/aruco4x4_100_",flip=False)
  #generate_tags(aruco.DICT_4X4_1000,"aruco_4x4_1000/aruco4x4_1000_",1000,flip=flip)
  #generate_tags(aruco.DICT_5X5_100,"aruco_5x5_100/aruco5x5_100_",flip=flip)
  #generate_tags(aruco.DICT_5X5_1000,"aruco_5x5_1000/aruco5x5_1000_",1000,flip=flip)
  print("complete!")
